import numpy as np
import mindspore as ms
import mindspore.nn as nn
from mindspore.ops import functional as F
from mindspore.common.initializer import initializer


class cnnEncoder(nn.Cell):
    def __init__(self):
        super(cnnEncoder, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1, pad_mode='valid', has_bias=True, weight_init='ones')
        self.conv2 = nn.Conv2d(20, 50, 5, 1, pad_mode='valid', has_bias=True, weight_init='ones')
        self.fc1 = nn.Dense(4*4*50, 500, weight_init='ones')
        self.fc2 = nn.Dense(500, 10, weight_init='ones')
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(2,2)

        
    def construct(self, x):
        x = self.conv1(x)
        x = self.relu(x)     
        x = self.max_pool2d(x)
        x = self.conv2(x)
        x = self.relu(x)    
        x = self.max_pool2d(x)
        x = x.view(-1, 4*4*50)      
        x = self.fc1(x)
        x = self.relu(x)   
        x = self.fc2(x)  

        return x

if __name__ == "__main__":
    encoder = cnnEncoder()
    weight_shape = encoder.fc.weight.shape
    weight = initializer('ones', shape = weight_shape, dtype=ms.float32)
    encoder.fc.weight = weight
    x = np.random.randn(1, 1, 28, 28).astype(np.float32)
    x = ms.Tensor(x, ms.float32)
    print(x)
    result = encoder(x)
    print(result)
    print(result.mean())
    